import pandas as pd
import baostock as bs
# 接口类,利用baostock获取数据
class get_stock:

    def __init__(self):

        self.name = '利用baostock接口获取A股数据'

    # 获取某一日期所有A股数据
    def get_all_stock_in_a_day(self,date):
        bs.login()
        print(date)

        # 获取指定日期的指数、股票数据
        stock_rs = bs.query_all_stock(date)
        stock_df = stock_rs.get_data()

        # 创建一个df文件储存数据
        data_df = pd.DataFrame()
        for code in stock_df["code"]:
            print("Downloading :" + code)
            k_rs = bs.query_history_k_data_plus(code, "date,code,open,high,low,close", date, date)
            data_df = data_df.append(k_rs.get_data())

        #data_df.append(stock_df["code_name"])
        # data_df["code_name"]=stock_df["code_name"]

        data_df.drop('date', axis=1, inplace=True)
        data_df.set_index('code', inplace=True)

        # 处理NaN：删除该行数据
        data_df = data_df.dropna()


        #### 登出系统 ####
        bs.logout()
        return data_df
        # data_df.to_csv("./stockData.csv", encoding="gbk", index=False)
        # print(data_df)

    # 利用bs获取交易日
    def get_trade_dates(self,start_date,end_date):
        #### 登陆系统 ####
        bs.login()
        #### 获取交易日信息 ####
        rs = bs.query_trade_dates(start_date=start_date, end_date=end_date)
        #### 打印结果集 ####
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(rs.get_row_data())
        result = pd.DataFrame(data_list, columns=rs.fields)
        trade = result[result.is_trading_day.isin(['1'])]
        date_list = list(trade['calendar_date'])

        #### 结果集输出到csv文件 ####
        #result.to_csv("D:\\trade_datas.csv", encoding="gbk", index=False)
        #print(result)

        #### 登出系统 ####
        bs.logout()

        return date_list

    # 获取某一阶段所有A股的数据，输出为字典{date:data_df}
    def get_all_stock(self,start_date,end_date):

        date_list=self.get_trade_dates(start_date,end_date)
        data={}
        for date in date_list:
            data_df=self.get_all_stock_in_a_day(date)
            data[date]=data_df
        return data

    # 获取一只股票的历史日线数据,输出为df
    def get_a_stock(self,stock_code,start_date,end_date):
        bs.login()
        rs = bs.query_history_k_data_plus(stock_code,
                                          "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                                          start_date=start_date, end_date=end_date,
                                          frequency="d", adjustflag="3")
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(rs.get_row_data())
        result = pd.DataFrame(data_list, columns=rs.fields)

        result.drop('code', axis=1, inplace=True)
        result.set_index('date', inplace=True)
        #### 登出系统 ####
        bs.logout()
        # 处理NAN：替换为0
        result.fillna(0)

        return result

    # 获取某一日期所有股票名称及其代码
    def all_stock(self,date):
        #### 登陆系统 ####
        bs.login()
        #### 获取证券信息 ####
        rs = bs.query_all_stock(day=date)
        #### 打印结果集 ####
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(rs.get_row_data())
        result = pd.DataFrame(data_list, columns=rs.fields)
        return result

        #### 登出系统 ####
        bs.logout()

    # 由股票代码或者名称(fuck)得到列表[名称，代码]
    def stock_code(self,fuck):
        result = pd.read_csv('all_stock.csv',encoding="utf-8")
        code = result.code.values
        name = result.code_name.values
        stocks1 = dict(zip(name, code))
        stocks2 = dict(zip(code, name))

        if '\u4e00' <= fuck[0] <= '\u9fff':
            fuckcode = stocks1[fuck]
            list = [fuck, fuckcode]
        else:
            fuckname = stocks2[fuck]
            list = [fuckname, fuck]
        return list



